package com.lnj.utils;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.ContentMetadata;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.stream.Collectors;

@Component
public class ProductContentRetriever implements ContentRetriever {

    private final EmbeddingStore<TextSegment> embeddingStore;
    private final EmbeddingModel embeddingModel;
    private static final int MAX_RESULTS = 5;

    @Autowired
    public ProductContentRetriever(EmbeddingStore<TextSegment> embeddingStore,
                                   EmbeddingModel embeddingModel) {
        this.embeddingStore = embeddingStore;
        this.embeddingModel = embeddingModel;
    }

    @Override
    public List<Content> retrieve(Query query) {
        System.out.println(query.text());
        // 1. 将查询文本向量化
        Embedding queryEmbedding = embeddingModel.embed(query.text()).content();

        // 2. 构建搜索请求
        EmbeddingSearchRequest request = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(MAX_RESULTS)
                .build();

        // 3. 执行搜索
        List<EmbeddingMatch<TextSegment>> relevantMatches = embeddingStore.search(request).matches();

        // 4. 转换为Content对象
        return relevantMatches.stream()
                .map(match -> {
                    // 从匹配的文本段中提取商品信息
                    String productInfo = match.embedded().text();
                    Content content = Content.from(productInfo);

                    // 添加元数据
                    content.metadata().put(ContentMetadata.valueOf("score"), match.score());
                    content.metadata().put(ContentMetadata.valueOf("productId"), extractProductId(productInfo));
                    return content;
                })
                .collect(Collectors.toList());
    }

    private Long extractProductId(String productInfo) {
        try {
            // 假设productInfo格式为："商品ID:123,名称:xxx,描述:xxx"
            String[] parts = productInfo.split(",");
            String idPart = parts[0].trim();
            return Long.parseLong(idPart.replace("商品ID:", "").trim());
        } catch (Exception e) {
            return null;
        }
    }
}